import time
from typing import Union

import cv2
import os
import gym
from numpy import average
import numpy as np

import ril_env
from ril.util import extract_min_rect, show_obj

BALL_R = [10, 14]
BAT_X = [140, 144]
ENEMY_X = [16, 20]
BAT_W = [6, 10]
BAT_H = [18, 22]
BAT_SPEED = 3
DO_NOTHING, MOVE_UP, MOVE_DOWN = 1, 2, 3
NOT_DETECTED_MAX = 20

ball_x, ball_y, bat_y, final_ball_y, enemy_y = None, None, None, None, None
not_detected_time = 0


def pong_rule(frame: Union[np.ndarray, None], act: int):
    """
    not_detected_time: 用于计数未检测到球的次数
    final_ball_y: 通过计算轨迹，球最后会到达的在球拍处的纵坐标
    """
    global ball_x, ball_y, bat_y, enemy_y, not_detected_time, final_ball_y
    slope, ball_detected = None, False
    if frame is None or frame.shape[-3:] != (210, 160, 3):
        return DO_NOTHING, None, False
    objects = extract_min_rect(frame)  # 查找物体边缘，得到每个物体的最小矩形
    for (index, (x, y, w, h)) in enumerate(objects):
        if BALL_R[0] < w < BALL_R[1] and BALL_R[0] < h < BALL_R[1]:
            if ball_x and ball_y:
                # 球运动轨迹的斜率
                slope = 0 if (x + w/2 - ball_x) == 0 else \
                    (y + h/2 - ball_y) / (x + w/2 - ball_x)
            ball_x, ball_y = x + w / 2, y + h / 2
            ball_detected = True
        elif BAT_X[0] < x + w / 2 < BAT_X[1] and BAT_W[0] < w < BAT_W[1]:
            bat_y = y + h / 2
        elif ENEMY_X[0] < x + w / 2 < ENEMY_X[1] and BAT_W[0] < w < BAT_W[1]:
            enemy_y = y + h / 2

    if bat_y is None:
        return DO_NOTHING, None, False

    if ball_detected and slope is not None:
        not_detected_time = 0
        final_ball_y = slope * (average(BAT_X) - ball_x) + ball_y
    else:
        not_detected_time += 1

    act = DO_NOTHING
    if final_ball_y is not None and not_detected_time < NOT_DETECTED_MAX:
        if final_ball_y > bat_y + BAT_SPEED:
            act = MOVE_DOWN
        elif final_ball_y < bat_y - BAT_SPEED:
            act = MOVE_UP
    elif enemy_y is not None:
        # 长时间检测不到球，可能是球位于上下边缘被遮挡了一部分
        # 这时我们跟着敌人运动
        if enemy_y > bat_y + BAT_SPEED:
            act = MOVE_DOWN
        elif enemy_y < bat_y - BAT_SPEED:
            act = MOVE_UP

    # 预测的球拍 y，这是为了防止下次检测不到球拍
    if act == MOVE_UP:
        bat_y -= BAT_SPEED
    elif act == MOVE_DOWN:
        bat_y += BAT_SPEED
    return act, None, False


if __name__ == '__main__':
    try:
        del os.environ['SDL_VIDEODRIVER']
    except KeyError:
        pass
    env = gym.make('PongNoFrameskip-v4')
    score = 0
    for _ in range(1):
        env.reset()
        observation = None
        while True:
            action, _, __ = pong_rule(observation, DO_NOTHING)
            observation, reward, done, info = env.step(action)
            if done:
                env.close()
                print('score: %d' % score)
                break
            score += reward
            env.render()
            # time.sleep(0.01)
